// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/container/f14_set.h>

#include <pollux/vector/type_aliases.h>
#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/checked_arithmetic.h>
#include <pollux/functions/lib/lambda_function_util.h>

namespace kumo::pollux::functions {
    namespace {
        ///
        /// Implements the array_sum function.
        /// See documentation at https://prestodb.io/docs/current/functions/array.html
        ///

        template<typename TInput, typename TOutput, bool kPropagateElementNull>
        class ArraySumFunction : public exec::VectorFunction {
        public:
            template<bool may_have_nulls, typename DataAtFunc, typename IsNullFunc>
            TOutput applyCore(
                vector_size_t row,
                const ArrayVector *array_vector,
                DataAtFunc &&dataAtFunc,
                IsNullFunc &&isNullFunc,
                bool &nullResult) const {
                auto start = array_vector->offsetAt(row);
                auto end = start + array_vector->sizeAt(row);
                TOutput sum = 0;

                auto addElement = [](TOutput &sum, TInput value) {
                    if constexpr (std::is_same_v<TOutput, int64_t>) {
                        sum = checkedPlus<TOutput>(sum, value);
                    } else {
                        sum += value;
                    }
                };

                for (auto i = start; i < end; i++) {
                    if constexpr (may_have_nulls) {
                        if (isNullFunc(i)) {
                            if constexpr (kPropagateElementNull) {
                                nullResult = true;
                                return 0;
                            }
                        } else {
                            addElement(sum, dataAtFunc(i));
                        }
                    } else {
                        addElement(sum, dataAtFunc(i));
                    }
                }
                if constexpr (may_have_nulls && kPropagateElementNull) {
                    nullResult = false;
                }
                return sum;
            }

            template<bool may_have_nulls, typename DataAtFunc, typename IsNullFunc>
            void applyCore(
                const SelectivityVector &rows,
                exec::EvalCtx &context,
                const ArrayVector *array_vector,
                FlatVector<TOutput> *resultValues,
                DataAtFunc &&dataAtFunc,
                IsNullFunc &&isNullFunc) const {
                context.applyToSelectedNoThrow(rows, [&](auto row) {
                    bool nullResult;
                    auto sum = applyCore<may_have_nulls>(
                        row, array_vector, dataAtFunc, isNullFunc, nullResult);
                    if constexpr (may_have_nulls && kPropagateElementNull) {
                        if (nullResult) {
                            resultValues->set_null(row, true);
                        } else {
                            resultValues->set(row, sum);
                        }
                    } else {
                        resultValues->set(row, sum);
                    }
                });
            }

            template<bool may_have_nulls>
            void applyFlat(
                const SelectivityVector &rows,
                exec::EvalCtx &context,
                ArrayVector *array_vector,
                const uint64_t *rawNulls,
                const TInput *rawElements,
                FlatVector<TOutput> *resultValues) const {
                applyCore<may_have_nulls>(
                    rows,
                    context,
                    array_vector,
                    resultValues,
                    [&](vector_size_t index) { return rawElements[index]; },
                    [&](vector_size_t index) { return bits::isBitNull(rawNulls, index); });
            }

            template<bool may_have_nulls>
            void applyNonFlat(
                const SelectivityVector &rows,
                exec::EvalCtx &context,
                ArrayVector *array_vector,
                exec::LocalDecodedVector &elements,
                FlatVector<TOutput> *resultValues) const {
                applyCore<may_have_nulls>(
                    rows,
                    context,
                    array_vector,
                    resultValues,
                    [&](vector_size_t index) {
                        return elements->template value_at<TInput>(index);
                    },
                    [&](vector_size_t index) { return elements->is_null_at(index); });
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args, // Not using const ref so we can reuse args
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                // Input is either flat or constant.

                if (args[0]->is_constant_encoding()) {
                    auto array_vector = args[0]->wrapped_vector()->as<ArrayVector>();
                    auto arrayRow = args[0]->wrapped_index(rows.begin());
                    POLLUX_CHECK(array_vector);
                    auto elementsVector = array_vector->elements();

                    SelectivityVector elementsRows(elementsVector->size());
                    exec::LocalDecodedVector elements(context, *elementsVector, elementsRows);

                    TOutput sum;
                    bool nullResult = false;
                    try {
                        if (elementsVector->may_have_nulls()) {
                            sum = applyCore<true>(
                                arrayRow,
                                array_vector,
                                [&](auto index) { return elements->value_at<TInput>(index); },
                                [&](auto index) { return elements->is_null_at(index); },
                                nullResult);
                        } else {
                            sum = applyCore<false>(
                                arrayRow,
                                array_vector,
                                [&](auto index) { return elements->value_at<TInput>(index); },
                                [&](auto index) { return elements->is_null_at(index); },
                                nullResult);
                        }
                    } catch (...) {
                        context.setErrors(rows, std::current_exception());
                    }

                    context.moveOrCopyResult(
                        std::make_shared<ConstantVector<TOutput> >(
                            context.pool(),
                            rows.end(),
                            nullResult,
                            outputType,
                            std::move(sum)),
                        rows,
                        result);
                    return;
                }

                POLLUX_CHECK_EQ(
                    args[0]->encoding(),
                    VectorEncoding::Simple::ARRAY,
                    "Expected flat or constant encoding");

                // Prepare result vector for writing
                BaseVector::ensure_writable(rows, outputType, context.pool(), result);
                auto resultValues = result->template as_flat_vector<TOutput>();

                auto array_vector = args[0]->as<ArrayVector>();
                POLLUX_CHECK(array_vector);
                auto elementsVector = array_vector->elements();

                if (elementsVector->is_flat_encoding()) {
                    const TInput *__restrict rawElements =
                            elementsVector->as<FlatVector<TInput> >()->rawValues();
                    const uint64_t *__restrict rawNulls = elementsVector->raw_nulls();

                    if (elementsVector->may_have_nulls()) {
                        applyFlat<true>(
                            rows, context, array_vector, rawNulls, rawElements, resultValues);
                    } else {
                        applyFlat<false>(
                            rows, context, array_vector, rawNulls, rawElements, resultValues);
                    }
                } else {
                    SelectivityVector elementsRows(elementsVector->size());
                    exec::LocalDecodedVector elements(context, *elementsVector, elementsRows);

                    if (elementsVector->may_have_nulls()) {
                        applyNonFlat<true>(rows, context, array_vector, elements, resultValues);
                    } else {
                        applyNonFlat<false>(rows, context, array_vector, elements, resultValues);
                    }
                }
            }
        };

        // Create function.
        template<bool kPropagateElementNull>
        std::shared_ptr<exec::VectorFunction> create(
            const std::string & /* name */,
            const std::vector<exec::VectorFunctionArg> &inputArgs,
            const core::QueryConfig & /*config*/) {
            auto elementType = inputArgs.front().type->childAt(0);

            switch (elementType->kind()) {
                case TypeKind::TINYINT: {
                    return std::make_shared<ArraySumFunction<
                        TypeTraits<TypeKind::TINYINT>::NativeType,
                        int64_t,
                        kPropagateElementNull> >();
                }
                case TypeKind::SMALLINT: {
                    return std::make_shared<ArraySumFunction<
                        TypeTraits<TypeKind::SMALLINT>::NativeType,
                        int64_t,
                        kPropagateElementNull> >();
                }
                case TypeKind::INTEGER: {
                    return std::make_shared<ArraySumFunction<
                        TypeTraits<TypeKind::INTEGER>::NativeType,
                        int64_t,
                        kPropagateElementNull> >();
                }
                case TypeKind::BIGINT: {
                    return std::make_shared<ArraySumFunction<
                        TypeTraits<TypeKind::BIGINT>::NativeType,
                        int64_t,
                        kPropagateElementNull> >();
                }
                case TypeKind::REAL: {
                    return std::make_shared<ArraySumFunction<
                        TypeTraits<TypeKind::REAL>::NativeType,
                        double,
                        kPropagateElementNull> >();
                }
                case TypeKind::DOUBLE: {
                    return std::make_shared<ArraySumFunction<
                        TypeTraits<TypeKind::DOUBLE>::NativeType,
                        double,
                        kPropagateElementNull> >();
                }
                default: {
                    POLLUX_FAIL("Unsupported Type");
                }
            }
        }

        // Define function signature.
        // array(T1) -> T2 where T1 must be coercible to bigint or double, and
        // T2 is bigint or double
        std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
            static const std::map<std::string, std::string> typePairs = {
                {"tinyint", "bigint"},
                {"smallint", "bigint"},
                {"integer", "bigint"},
                {"bigint", "bigint"},
                {"real", "double"},
                {"double", "double"}
            };
            std::vector<std::shared_ptr<exec::FunctionSignature> > signatures;
            signatures.reserve(typePairs.size());
            for (const auto &[argType, returnType]: typePairs) {
                signatures.emplace_back(exec::FunctionSignatureBuilder()
                    .returnType(returnType)
                    .argumentType(fmt::format("array({})", argType))
                    .build());
            }
            return signatures;
        }
    } // namespace

    // Register function.
    void registerVectorFunction_udf_array_sum(const std::string &name) {
        kumo::pollux::exec::registerStatefulVectorFunction(
            name, signatures(), create<false>);
        kumo::pollux::exec::registerStatefulVectorFunction(
            name + "_propagate_element_null", signatures(), create<true>);
    }
} // namespace kumo::pollux::functions
